# 从路径中分离目录部分和文件名部分
from os.path import split
# 用于解析命令行参数
import argparse
# 提供日志功能，用于记录程序运行时的信息。
import logging
# 提供操作系统相关的功能，例如路径操作、文件操作等
import os
# 生成随机数，设置随机种子
import random
# 与 Python 解释器交互
import sys
# 用于科学计算，支持多维数组操作、数学函数等。
import numpy as np
# 核心深度学习框架，提供张量运算、自动求导和神经网络构建功能
import torch
# CUDNN 是 NVIDIA 提供的 GPU 加速库，PyTorch 提供相关的接口。
import torch.backends.cudnn as cudnn
# 数据加载器，用于将数据集分批次加载到模型中
from torch.utils.data import DataLoader
# 显示进度条的工具。
from tqdm import tqdm
# tensorboard初始化
from torch.utils.tensorboard import SummaryWriter
# 获取程序的配置
from config import get_config
# 自定义数据集类，用于处理 Synapse 数据集。
from datasets.dataset_synapse import Synapse_dataset
# 实现 Swin Transformer 和 U-Net 的组合模型。
from networks.vision_transformer import SwinUnet as ViT_seg
from networks.unetPP import UNetPP as UnetPP
from networks.unet import UNet as UNet
from networks.u2net import U2Net as U2Net
from networks.vgg_unet import VGG16UNet as VGG16UNet
from networks.mobilenet_unet import MobileV3Unet as MobileUnet
# 验证阶段单次推理的工具函数
from utils import test_single_volume

parser = argparse.ArgumentParser()
# 指定数据的根路径，通常用于验证集或测试集
parser.add_argument('--root_path', type=str,
                    default='./datasets/Synapse/test_vol_h5',
                    help='root dir for validation volume data')  # acdc volume_path=root_dir
# 数据集的列表文件目录。
parser.add_argument('--list_dir', type=str,
                    default='./lists/Synapse', help='list dir')
# 用于指定当前运行的实验数据集，例如 Synapse
parser.add_argument('--dataset', type=str,
                    default='datasets', help='experiment_name')
# 网络的输出通道数（即分类数）
parser.add_argument('--num_classes', type=int,
                    default=9, help='output channel of network')
# 已经训练好的模型加载路径
# 尝试加载最佳模型 best_model.pth
# 如果不存在，则加载最后一个 epoch 的权重。
parser.add_argument('--output_dir', type=str, help='output dir')
# 最大迭代次数
parser.add_argument('--max_iterations', type=int, default=30000, help='maximum epoch number to train')
# 最大训练 epoch 数
parser.add_argument('--max_epochs', type=int, default=150, help='maximum epoch number to train')
# 每个 GPU 的批次大小
parser.add_argument('--batch_size', type=int, default=24,
                    help='batch_size per gpu')
# 网络输入图像的尺寸
parser.add_argument('--img_size', type=int, default=224, help='input patch size of network input')
# 是否在推理过程中保存预测结果为 .nii 格式
# action="store_true" 的作用是，如果在命令行中指定了 --is_savenii 参数，那么 args.is_savenii 会被设置为 True，否则默认为 False。
parser.add_argument('--is_savenii', action="store_true", help='whether to save results during inference')
# 推理结果保存路径，默认：../predictions
parser.add_argument('--test_save_dir', type=str, default='../predictions', help='saving prediction as nii!')
# 是否启用确定性训练
parser.add_argument('--deterministic', type=int, default=1, help='whether use deterministic training')
# 基础学习率
parser.add_argument('--base_lr', type=float, default=0.01, help='segmentation network learning rate')
# 随机数种子
parser.add_argument('--seed', type=int, default=3407, help='random seed')
# 配置文件路径
parser.add_argument('--cfg', type=str, required=True, metavar="FILE", help='path to config file', )
# 用于动态修改配置选项（以键值对形式传递）
parser.add_argument(
    "--opts",
    help="Modify config options by adding 'KEY VALUE' pairs. ",
    default=None,
    nargs='+',
)
# 是否使用压缩数据集
parser.add_argument('--zip', action='store_true', help='use zipped dataset instead of folder dataset')
# 数据加载时的缓存模式
parser.add_argument('--cache-mode', type=str, default='part', choices=['no', 'full', 'part'],
                    help='no: no cache, '
                         'full: cache all data, '
                         'part: sharding the dataset into nonoverlapping pieces and only cache one piece')
# 从检查点恢复训练
parser.add_argument('--resume', help='resume from checkpoint')
# 梯度累积步数（用于在显存不足时模拟更大的 batch size）
parser.add_argument('--accumulation-steps', type=int, help="gradient accumulation steps")
# 是否使用检查点以节省显存
parser.add_argument('--use-checkpoint', action='store_true',
                    help="whether to use gradient checkpointing to save memory")
# 混合精度的优化等级。
parser.add_argument('--amp-opt-level', type=str, default='O1', choices=['O0', 'O1', 'O2'],
                    help='mixed precision opt level, if O0, no amp is used')
# 实验标签，用于标记特定实验
parser.add_argument('--tag', help='tag of experiment')
# 是否仅执行验证
parser.add_argument('--eval', action='store_true', help='Perform evaluation only')
# 是否仅测试模型的推理吞吐量
parser.add_argument('--throughput', action='store_true', help='Test throughput only')
# 存放验证集的列表文件名，默认值为 test
parser.add_argument("--split_name", default="test", help="Directory of the input list")

# 保存文件
args = parser.parse_args()
config = get_config(args)
# 验证
def inference(args, model, test_save_path=None):
    # 使用自定义数据集类 Synapse_dataset 加载测试数据集
    # base_dir：数据集根目录，通常指定测试数据所在路径
    # split：数据集的分割名称，例如 test
    # list_dir：数据列表文件所在的目录
    # print(args.list_dir)
    db_test = Synapse_dataset(base_dir=args.root_path, split=args.split_name, list_dir=args.list_dir)
    # 数据加载器，使用 PyTorch 的 DataLoader
    # 批次大小设置为 1
    # 数据顺序不打乱（shuffle=False）
    # 数据操作线程设置为 1
    testloader = DataLoader(db_test, batch_size=1, shuffle=False, num_workers=1)
    logging.info("一轮的测试迭代数为{}".format(len(testloader)))
    # 模型切换到验证模式
    model.eval()
    # 初始化性能指标（Dice 系数和 Hausdorff 距离）
    metric_list = 0.0
    # 遍历测试数据集
    # 使用 tqdm 显示进度条，方便观察测试进度
    with torch.no_grad():
        for i_batch, sampled_batch in tqdm(enumerate(testloader)):
            # h, w = sampled_batch["image"].size()[2:]
            # 从 sampled_batch 中提取数据：
            # image：输入图像。
            # label：真实标签。
            # case_name：测试样本的名称，便于保存结果或记录日志。
            image, label, case_name = sampled_batch["image"], sampled_batch["label"], sampled_batch['case_name'][0]
            # 处理数据集名称（可选）
            if args.dataset == "datasets":
                case_name = split(case_name.split(",")[0])[-1]
            # 调用工具函数 test_single_volume，对单个样本进行推理，并计算性能指标。
            # 输入参数：
                # image 和 label：单张图像及其对应标签。
                # model：要测试的深度学习模型。
                # classes：类别数量（args.num_classes）。
                # patch_size：模型输入图像的大小（[args.img_size, args.img_size]）。
                # test_save_path：测试结果的保存路径。
                # case：当前样本的名称。
                # z_spacing：3D 图像的层间距离。
            # 返回值：
                # metric_i：当前样本的指标，例如 Dice 系数和 Hausdorff 距离。
            metric_i = test_single_volume(image, label, model, classes=args.num_classes,
                                        patch_size=[args.img_size, args.img_size],
                                        test_save_path=test_save_path, case=case_name, z_spacing=args.z_spacing)
            # 累积指标
            metric_list += np.array(metric_i)
            # 日志记录单样本性能
            logging.info('idx %d case %s mean_dice %f mean_hd95 %f' % (
                i_batch, case_name, np.mean(metric_i, axis=0)[0], np.mean(metric_i, axis=0)[1]))
            
    # 计算整体指标
    metric_list = metric_list / len(db_test)
    # 按类别记录指标
    for i in range(1, args.num_classes):
        logging.info('Mean class %d mean_dice %f mean_hd95 %f' % (i, metric_list[i - 1][0], metric_list[i - 1][1]))
    # 记录整体性能
    performance = np.mean(metric_list, axis=0)[0]
    mean_hd95 = np.mean(metric_list, axis=0)[1]
    logging.info('Testing performance in best val model: mean_dice : %f mean_hd95 : %f' % (performance, mean_hd95))
    # 返回测试结果
    return "Testing Finished!"


if __name__ == "__main__":
    # 设置是否使用确定性训练（即结果可复现）
    if not args.deterministic:
        # 启用 cuDNN 的加速优化，但可能导致结果不稳定
        cudnn.benchmark = True
        # 启用确定性模式，确保每次运行结果一致
        cudnn.deterministic = False
    else:
        cudnn.benchmark = False
        cudnn.deterministic = True

    # 设置随机种子，确保随机数生成的可重复性
    random.seed(args.seed)
    np.random.seed(args.seed)
    torch.manual_seed(args.seed)
    torch.cuda.manual_seed(args.seed)

    # 初始化模型
    args.is_pretrain = True
    args.z_spacing = 1
    net = ViT_seg(config, img_size=args.img_size, num_classes=args.num_classes).cuda()
    # net = UnetPP().cuda()
    # net = UNet().cuda()
    # net = U2Net().cuda()
    # net = VGG16UNet().cuda()
    # net = MobileUnet(num_classes=9).cuda()

    # 加载模型权重：尝试加载最佳模型 best_model.pth，如果不存在，则加载最后一个 epoch 的权重。
    snapshot = os.path.join(args.output_dir, 'best_model.pth')
    if not os.path.exists(snapshot):
        snapshot = snapshot.replace('best_model', 'epoch_' + str(args.max_epochs - 1))
    msg = net.load_state_dict(torch.load(snapshot))
    print("win unet预训练权重如下:", msg)
    snapshot_name = snapshot.split('/')[-1]

    # 日志记录：logging 用于记录测试过程中的关键信息。
    log_folder = args.output_dir +'/test_log/test_log_'
    writer = SummaryWriter(log_dir=log_folder) 

    os.makedirs(log_folder, exist_ok=True)
    logging.basicConfig(filename=log_folder + '/' + snapshot_name + ".txt", level=logging.INFO,
                        format='[%(asctime)s.%(msecs)03d] %(message)s', datefmt='%H:%M:%S')
    logging.getLogger().addHandler(logging.StreamHandler(sys.stdout))
    logging.info(str(args))
    logging.info(snapshot_name)
    # 配置结果保存路径
    # 如果设置了 args.is_savenii，保存预测结果到 predictions 文件夹。
    if args.is_savenii:
        args.test_save_dir = os.path.join(args.output_dir, "predictions")
        test_save_path = args.test_save_dir
        os.makedirs(test_save_path, exist_ok=True)
    # 否则不保存预测结果。
    else:
        test_save_path = None
    # 执行推理
    inference(args, net, test_save_path)

# python train.py --dataset Synapse --cfg $CFG --root_path $DATA_DIR --max_epochs $EPOCH_TIME --output_dir $OUT_DIR --img_size $IMG_SIZE --base_lr $LEARNING_RATE --batch_size $BATCH_SIZE
# python train.py --output_dir './model_out/datasets' --dataset datasets --img_size 224 --batch_size 32 --cfg configs/swin_tiny_patch4_window7_224_lite.yaml --root_path /media/aicvi/11111bdb-a0c7-4342-9791-36af7eb70fc0/NNUNET_OUTPUT/nnunet_preprocessed/Dataset001_mm/nnUNetPlans_2d_split
# python test.py --output_dir ./model_out/datasets --dataset datasets --cfg configs/swin_tiny_patch4_window7_224_lite.yaml --is_saveni --root_path /media/aicvi/11111bdb-a0c7-4342-9791-36af7eb70fc0/NNUNET_OUTPUT/nnunet_preprocessed/Dataset001_mm/test --max_epoch 150 --base_lr 0.05 --img_size 224 --batch_size 24
